
# coding: utf-8

# ## 导入各种依赖库
# 因为此目录下已经有训练好的模型，可以直接到最后一个代码块运行，无需额外训练，如果与实际应用情况存在偏差，请添加数据集进行再次训练
# 
# 没有的模块进行安装
# ``` bash
# pip install modual
# ```

# In[2]:


import tensorflow as tf
import gzip
import os
import tempfile
import cv2
import random
import numpy as np


# ## 生成数据集格式

# ### train data

# In[3]:


#导入数据图片，以features命名
imgs=os.listdir('./datasets/train/imgs/')
# 定义一个排序函数
def nu_str(string):
    return int(string.split('.')[0])
# 将文件夹中的文件按照名称数字大小进行排序 能够与labels一一对应
imgs.sort(key=nu_str)
features_train=[]
# 对每一张图片进行处理，主要是将矩阵转化为一个向量，最后将所有图片打包
for i in imgs:
    img=cv2.imread('./datasets/train/imgs/'+str(i),0)
    #res,img=cv2.threshold(img, 200, 255, cv2.THRESH_BINARY)
    #img=cv2.copyMakeBorder(img,5,5,5,5,cv2.BORDER_CONSTANT,value=0)
    #cv2.imshow('3',img)
    #cv2.waitKey(100)
    img=img.reshape(28*28)/255
    features_train.append(img)
features_train=np.array(features_train) # 包含所有图片的一个向量集


# In[4]:


## 将每一个图片对应的结果转化为one-hot形式储存
# 读取文件所有内容
with open('./datasets/train/targets/target.txt','r') as f:
    tars=f.readlines()
# 向量不同位置对应的结果
tar_temp=[0,1,2,3,4,5,6,7,8,9,'.']
labels_train=[]
# 构造one-hot形式的向量集
for i in tars:
    b=np.array([i[0]==str(tar_temp[j]) for j in range(len(tar_temp))])+0
    labels_train.append(b)  # 一个包含所有结果的向量集（与图片集一一对应）


# ### 验证数据集的正确性

# In[5]:


# 选择运行
# 查看数据集与结果是否一一对应，主要看看显示的图片和打印的数字是否一致

for i in range(len(features_train)):
    cv2.imshow('feature',features_train[i].reshape(28,28))
    print(np.argmax(labels_train[i]))
    cv2.waitKey(500) # 单张图片的显示时间ms


# ### test data

# In[6]:


# 生成测试数据集，同训练集一样

imgs=os.listdir('./datasets/test/imgs/')
def nu_str(string):
    return int(string.split('.')[0])
imgs.sort(key=nu_str)
features_test=[]
for i in imgs:
    img=cv2.imread('./datasets/test/imgs/'+str(i),0)
    #res,img=cv2.threshold(img, 200, 255, cv2.THRESH_BINARY)
    #img=cv2.copyMakeBorder(img,5,5,5,5,cv2.BORDER_CONSTANT,value=0)
    #cv2.imshow('3',img)
    #cv2.waitKey(100)
    img=img.reshape(28*28)/255
    features_test.append(img)
features_test=np.array(features_test)


# In[7]:


with open('./datasets/test/targets/target.txt','r') as f:
    tars=f.readlines()
tar_temp=[0,1,2,3,4,5,6,7,8,9,'.']
labels_test=[]
for i in tars:
    b=np.array([i[0]==str(tar_temp[j]) for j in range(len(tar_temp))])+0
    labels_test.append(b)


# In[8]:


# 查看测试训练集的正确性  可以跳过
for i in range(len(features_test)):
    cv2.imshow('feature',features_test[i].reshape(28,28))
    print(np.argmax(labels_test[i]))
    cv2.waitKey(500)


# ## 线性训练
# 这个训练结果太差，忽略

# ### train 
# 忽略 不需要运行

# In[9]:


sess = tf.Session()
# 设置放置自变量的空间(图片输入的)
x=tf.placeholder("float",[None,28*28])
# 设置权重以及偏置的变量空间
W=tf.Variable(tf.zeros([28*28,11]))
b=tf.Variable(tf.zeros([11]))
# 计算特定权重偏置下的结果
y=tf.nn.softmax(tf.matmul(x,W)+b)
# 设置储存实际结果的空间
y_=tf.placeholder("float",[None,11])
# 计算预测与实际的交叉熵
cross_entropy=-tf.reduce_sum(y_*tf.log(y))
# 训练 通过使得交叉熵向减小的方向
train_step=tf.train.GradientDescentOptimizer(0.01).minimize(cross_entropy)
# 初始化变量
init = tf.initialize_all_variables()
# 设置学习的session
sess.run(init)

#设置训练
for i in range(10):
    # 每次训练随机取50张图片和结果进行训练
    sample=random.sample(range(len(labels_train)),50)
    # 读取训练图片
    batch_xs=np.array([features_train[i] for i in sample])
    # 读取训练结果labels
    batch_ys=np.array([labels_train[i] for i in sample])
    # 训练
    sess.run(train_step, feed_dict={x: batch_xs, y_: batch_ys})


# ### test
# 不需要运行

# In[10]:


correct_prediction = tf.equal(tf.argmax(y,1), tf.argmax(y_,1))  # 比较结果和实际结果
accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float")) # 正确率计算
print(sess.run(accuracy, feed_dict={x: features_test, y_: labels_test})) # 测试
sess.close() # 关闭sess


# ## 多重卷积训练
# 正确率极高推荐使用

# In[11]:


# 定义权重函数工厂函数（批量生产权重的函数，为了方便）
def weight_variable(shape,name):
    initial = tf.truncated_normal(shape, stddev=0.1)
    return tf.Variable(initial,name=name)
# 定义偏置工厂函数
def bias_variable(shape,name):
    initial = tf.constant(0.1, shape=shape)
    return tf.Variable(initial,name=name)
# 定义卷积矩阵工厂函数
def conv2d(x, W):
      return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')
# 定义池化层矩阵工厂函数
def max_pool_2x2(x):
      return tf.nn.max_pool(x, ksize=[1, 2, 2, 1],
                        strides=[1, 2, 2, 1], padding='SAME')


# In[12]:


## 主要有两层卷积运算
# 第一层卷积层定义
W_conv1 = weight_variable([5, 5, 1, 32],name='w_conv1')  # 权重变量
b_conv1 = bias_variable([32],name='b_conv1',)            # 偏置

# 图片输入空间生成,结果空间生成（请求组织先分配好茅坑）
x = tf.placeholder("float", shape=[None, 28*28],name="X")  # 
y_ = tf.placeholder("float", shape=[None, 11],name="Y")

# 将输入空间重新塑造为28*28*1（1指单通道，-1是指可以随机应变），为了后面的卷积运算 因为输入是一个向量集
x_image = tf.reshape(x, [-1,28,28,1]) 

# 定义卷积矩阵并计算
h_conv1 = tf.nn.relu(conv2d(x_image, W_conv1) + b_conv1)
# 定义池化层 
h_pool1 = max_pool_2x2(h_conv1)
# 第二层卷积层定义
W_conv2 = weight_variable([5, 5, 32, 64],name='w_conv2')  # 权重变量
b_conv2 = bias_variable([64],name='b_conv2')              # 偏置

h_conv2 = tf.nn.relu(conv2d(h_pool1, W_conv2) + b_conv2)
h_pool2 = max_pool_2x2(h_conv2)

# 经过两次卷积和池化，最后的图片只有7*7了，但是还是不知道他到底是什么鬼，所以再来一个权重矩阵，来算算他到底是什么鬼
# 第一个与处理后图片尺寸一样的权重矩阵变量和偏置变量，直接点乘
W_fc1 = weight_variable([7 * 7 * 64, 1024],name='w_fc1')  
b_fc1 = bias_variable([1024],name='b_fc1')

h_pool2_flat = tf.reshape(h_pool2, [-1,7*7*64])
h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, W_fc1) + b_fc1)

# keep_prob为了防止过拟合，具体原理我还没看到。。。。
keep_prob = tf.placeholder("float",name='keep_prob')
h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)
# 再来一个矩阵使得其变成一个全连接层，所谓全连接层就是一个向量，之所以要将矩阵化为全连接层
# 就是为了使得他通过和再一个权重相乘能够得到和结果维度相同的输出
W_fc2 = weight_variable([1024, 11],name='w_fc2')
b_fc2 = bias_variable([11],name='b_fc2')
# 最后的结果输出为一个向量 和labels相同的维度，
y_conv=tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2)
# 设置模型格式 ，添加输出的格式进去
tf.add_to_collection('yconv',y_conv)
saver = tf.train.Saver()


# In[13]:


# 训练模型
with tf.Session() as sess:
    # 设置交叉熵为损失函数
    cross_entropy = -tf.reduce_sum(y_*tf.log(y_conv))
    # 设置优化参数，采用AdamOptimizer优化方法，比最速下降法更优，能够防止过拟合
    train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)
    # 判断预测结果和真实结果是否相同
    correct_prediction = tf.equal(tf.argmax(y_conv,1), tf.argmax(y_,1))
    # 精度
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))
    # 初始化各个变量
    sess.run(tf.initialize_all_variables())
    # 迭代训练
    for i in range(201):
        # 随机选取数据进行训练
        sample = random.sample(range(len(labels_train)),50)
        batch_xs=np.array([features_train[i] for i in sample])
        batch_ys=np.array([labels_train[i] for i in sample])
        # 当是100倍数是保存模型，并且输出当前测试精度，保存路径为相对路径
        if i%100 == 0:
            train_accuracy = accuracy.eval(feed_dict={x:batch_xs, y_: batch_ys, keep_prob: 1.0})
            print ("step %d, training accuracy %g"%(i, train_accuracy))
            save_path = saver.save(sess, "./datasets/digit_model/my_digit_model")
        train_step.run(feed_dict={x:batch_xs, y_: batch_ys, keep_prob: 0.5})
    # 测试整体精度，加载测试集
    print ("test accuracy %g"%accuracy.eval(feed_dict={x: features_test, y_: labels_test, keep_prob: 1.0}))


# ## 加载使用模型  应用训练好的模型
# 模型已经存在相应位置，可以直接使用此代码块

# In[1]:


import cv2
import numpy as np
import tensorflow as tf
import os



# 不支持多行数字识别，以及单行多个小数点的数值识别（单行只能实现字符串识别），
# labels的各个位置代表的数字
tar_temp=[0,1,2,3,4,5,6,7,8,9,'.'] 

# 定义一个阈值函数，将数码管部分取出来，根据实际情况进行相应修改，找到最优参数
def thresholding_inv(image):
    # 定义膨胀核心，根据实际情况进行修改
    kernel_dilate = cv2.getStructuringElement(cv2.MORPH_RECT,(1, 6))# 1代表横向膨胀，6代表纵向膨胀
    ## 腐蚀参数我已经注释掉，根据实际情况选择是否使用
    kernel_erode = cv2.getStructuringElement(cv2.MORPH_RECT,(2, 1)) 
    ## 根据RGB图得到灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # 灰度图二值化
    ret, bin = cv2.threshold(gray, 220, 255, cv2.THRESH_BINARY)
    ## 对灰度图进行腐蚀，主要是为了分离相近的小数点，如果足够清晰可以不使用腐蚀，我已注释掉
    bin = cv2.erode(bin,kernel_erode)
    ## 对灰度图进行膨胀
    bin=cv2.dilate(bin,kernel_dilate,iterations = 1)
    return bin

# Read the input image
## demo 图像在此目录下
im = cv2.imread('./datasets/2018-09-10_162811_270.jpg')  # 还有 1-6 张图 修改最后一个数即可
## 二值化处理
im_th = thresholding_inv(im)

# 显示图片
cv2.imshow('im_th',im_th)
cv2.waitKey(10) # 显示1000ms 

# Find contours in the image  寻找边界集合
_,ctrs, hier = cv2.findContours(im_th.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# Get rectangles contains each contour 
rects = [cv2.boundingRect(ctr) for ctr in ctrs]
# 加载训练好的模型，并预测通过
with tf.Session() as sess:
    # 加载模型的结构框架graph
    new_saver = tf.train.import_meta_graph('./datasets/digit_model/my_digit_model.meta')
    # 加载各种变量
    new_saver.restore(sess,'./datasets/digit_model/my_digit_model')
    yy_hyp = tf.get_collection('yconv')[0]
    graph = tf.get_default_graph() 
    X = graph.get_operation_by_name('X').outputs[0]#为了将 x placeholder加载出来
    keep_prob = graph.get_operation_by_name('keep_prob').outputs[0] # 将keep_prob placeholder加载出来
    # mm用来保存数字以及数字坐标
    mm={}
    # for循环对每一个contour 进行预测和求解，并储存
    for rect in rects:
        # Draw the rectangles 得到数字区域 roi
        cv2.rectangle(im, (rect[0], rect[1]), (rect[0] + rect[2], rect[1] + rect[3]), (0, 255, 0), 3) 
        # Make the rectangular region around the digit
        leng1= int(rect[3])
        leng2= int(rect[2])
        pt1 = int(rect[1] )
        pt2 = int(rect[0] )
        # 得到数字区域
        roi = im_th[pt1:pt1+leng1, pt2:pt2+leng2]
        # 尺寸缩放为模型尺寸
        roi = cv2.resize(roi, (28, 28), interpolation=cv2.INTER_AREA)
        # 处理成一个向量，为了和模型输入一直
        roi=np.array([roi.reshape(28*28)/255])
        # 运行模型得到预测结果
        pred= sess.run(yy_hyp,feed_dict = {X:roi,keep_prob:1.0})
        # 得到最大可能值索引 ind
        ind=np.argmax(pred)
        #labels不同位置代表的不同数字   (tar_temp[ind]) 就是预测值
        # 将预测值添加到图像中，并显示
        cv2.putText(im, str(tar_temp[ind]), (rect[0], rect[1]),cv2.FONT_HERSHEY_DUPLEX, 2, (0, 255, 255), 3)
        # 储存每个数字和其对应的boundingbox的像素点坐标
        mm[pt2]=tar_temp[ind]
    # 最后的处理
    # 根据像素坐标，从左到右排序，得到数字的顺序
    num_tup=sorted(mm.items(),key=lambda x:x[0])
    # 将数字列表连接为字符串
    num=(''.join([str(i[1]) for i in num_tup]))
    try:
        numn=float(num)
        print('图中数字为%s,数值大小为%s' %(num,numn))
    except:
        print('不好意思，目前不支持多个小数点的数值识别')
        print('图中数字为%s'% num)
    # 显示图像 
    cv2.namedWindow("Resulting Image with Rectangular ROIs", cv2.WINDOW_NORMAL)
    cv2.imshow("Resulting Image with Rectangular ROIs", im)
    cv2.waitKey(10)

